#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模型工厂模块
"""

import logging
import torch

from .spatiotemporal_fusion import SpatioTemporalFusion
from .i3d import I3D
from .video_mae import VideoMAE
from .swin_transformer_3d import SwinTransformer3D


def create_model(config):
    """
    创建模型
    
    Args:
        config (dict): 配置字典
    
    Returns:
        torch.nn.Module: 创建的模型
    """
    logger = logging.getLogger()
    
    # 获取模型配置
    model_config = config.get('model', {})
    model_type = model_config.get('type', 'SpatioTemporalFusion')
    device = config.get('system', {}).get('device', 'cuda' if torch.cuda.is_available() else 'cpu')
    
    logger.info(f"创建模型: {model_type}")
    
    # 根据模型类型创建对应的模型
    if model_type == 'SpatioTemporalFusion':
        model = SpatioTemporalFusion(
            backbone=model_config.get('backbone', 'resnet50'),
            pretrained=model_config.get('pretrained', True),
            feature_dim=model_config.get('feature_dim', 512),
            num_classes=model_config.get('num_classes', 7),
            dropout=model_config.get('dropout', 0.5),
            use_attention=model_config.get('use_attention', True),
            fusion_method=model_config.get('fusion_method', 'concat'),
            spatial_encoder=model_config.get('spatial_encoder', {}),
            temporal_encoder=model_config.get('temporal_encoder', {}),
            optical_flow_encoder=model_config.get('optical_flow_encoder', {}),
            num_frames=config.get('dataset', {}).get('num_frames', 16),
            frame_size=config.get('dataset', {}).get('frame_size', [224, 224])
        )
    
    elif model_type == 'I3D':
        model = I3D(
            num_classes=model_config.get('num_classes', 7),
            pretrained=model_config.get('pretrained', True),
            dropout_prob=model_config.get('dropout', 0.5)
        )
    
    elif model_type == 'VideoMAE':
        model = VideoMAE(
            num_classes=model_config.get('num_classes', 7),
            pretrained=model_config.get('pretrained', True),
            feature_dim=model_config.get('feature_dim', 768),
            dropout=model_config.get('dropout', 0.5)
        )
    
    elif model_type == 'SwinTransformer3D':
        model = SwinTransformer3D(
            num_classes=model_config.get('num_classes', 7),
            pretrained=model_config.get('pretrained', True),
            feature_dim=model_config.get('feature_dim', 768),
            dropout=model_config.get('dropout', 0.5)
        )
    
    else:
        logger.error(f"不支持的模型类型: {model_type}")
        raise ValueError(f"不支持的模型类型: {model_type}")
    
    # 将模型移动到指定设备
    model = model.to(device)
    
    # 打印模型参数数量
    num_params = sum(p.numel() for p in model.parameters())
    logger.info(f"模型参数数量: {num_params:,}")
    
    return model 